#codeing=utf-8
"""
    Candidates detection:

    date: 2017.04.04
"""



from easydict import EasyDict as edict
import os
import sys
import pprint
import argparse
import time
import subprocess

# os.environ["GLOG_minloglevel"] = "2"

from tools.dj_test_net_vol_luna_kaggle import test_process
from tools.cdd_use_lungmask import cdd_filter_lungmask
from tools.gen_candis_list import gen_list_one_fold


def parse_args():
    """
    Parse input arguments
    """
    parser = argparse.ArgumentParser(description='Detect Candidates')
    parser.add_argument('--gpu', dest='gpu_id', help='GPU id to use',
                        default=0, type=int)
    parser.add_argument('--start_idx', dest='start_scan_idx',
                        help='start_scan_idx',
                        default=None, type=int)
    parser.add_argument('--end_idx', dest='end_scan_idx',
                        help='end_scan_idx',
                        default=None, type=int)

    args = parser.parse_args()
    return args


if __name__ == '__main__':
    log_file_name = time.strftime('log/predict_%Y-%m-%d_%H-%M-%S.txt', time.localtime(time.time()))
    sys.stdout = os.fdopen(sys.stdout.fileno(), 'w', 0)
    tee = subprocess.Popen(["tee", log_file_name], stdin=subprocess.PIPE)
    os.dup2(tee.stdin.fileno(), sys.stdout.fileno())
    os.dup2(tee.stdin.fileno(), sys.stderr.fileno())

    cfg = edict()

    args = parse_args()
    print('Called with args:')
    print(args)

    if args.gpu_id is not None:
        cfg.gpuid = args.gpu_id
    else:
        cfg.gpuid = 0
    if args.start_scan_idx is not None:
        cfg.start_scan_idx = args.start_scan_idx
    else:
        cfg.start_scan_idx = 0
    if args.end_scan_idx is not None:
        cfg.end_scan_idx = args.end_scan_idx
    else:
        cfg.end_scan_idx = -1

    cfg.max_iter = 90000
    cwd = os.getcwd()
    cfg.data_root_dir = cwd + '/../../../data/'
    cfg.caffemodel_path = cfg.data_root_dir + 'tmp_faster_rcnn/v9/caffemodel'
    cfg.candi_dir = cfg.data_root_dir + 'tmp_faster_rcnn/v9/candidates'
    cfg.fold_list_path = cwd + '/../fold_list_all.npy'
    cfg.kaggle_vol_path = None
    cfg.lidc_vol_path = None
    cfg.spie_vol_path = None
    cfg.stage2_vol_path = cfg.data_root_dir + 'stage2/vol.hdf5'  # config

    cfg.candis_thresh = 0.98
    cfg.kaggle_lungmask_path = None
    cfg.lidc_lungmask_path = None
    cfg.spie_lungmask_path = None
    cfg.stage2_lungmask_path = cfg.data_root_dir + 'stage2/lungmask.hdf5'  # config
    cfg.candis_lungmask_dir = cfg.data_root_dir + 'tmp_faster_rcnn/v9/candidates_lungmask'
    cfg.pkl_dir = cfg.data_root_dir + 'faster_rcnn_candidates'

    if not os.path.exists(cfg.pkl_dir):
        os.makedirs(cfg.pkl_dir)

    fold_names = {1: 'fold1',
                  2: 'fold2',
                  3: 'fold3',
                  4: 'fold4',
                  5: 'fold5',
                  6: 'kaggle_beni',
                  7: 'kaggle_testset',
                  8: 'spie',
                  9: 'stage2'}

    print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
    print('------------------------- Testing --------------------------')
    print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
    fold_idx = 9
    fold_name = fold_names[fold_idx]
    print('----------------------- Testing ' + fold_name + ' ---------------------')
    caffemodel_name = os.path.join(
        cfg.caffemodel_path, 'fold1',
        'vgg16_faster_rcnn_iter_' + str(cfg.max_iter) + '.caffemodel')
    test_process(gpu_id=cfg.gpuid,
                 caffemodel=caffemodel_name,
                 candi_dir=os.path.join(cfg.candi_dir, fold_name),
                 fold_idx=fold_idx,  # 9 for kaggle stage2 test data
                 start_scan_idx=cfg.start_scan_idx,
                 end_scan_idx=cfg.end_scan_idx,
                 fold_list_path=cfg.fold_list_path,
                 kaggle_vol_path=cfg.kaggle_vol_path,
                 lidc_vol_path=cfg.lidc_vol_path,
                 spie_vol_path=cfg.spie_vol_path,
                 stage2_vol_path=cfg.stage2_vol_path,
                 prototxt='models/pascal_voc/VGG16/faster_rcnn_end2end/test.prototxt',
                 flag_wait=True
                 )

    print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
    print('------------------------- Lung Mask ------------------------')
    print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
    fold_idx = 9
    fold_name = fold_names[fold_idx]
    print('----------------------- LungMask ' + fold_name + ' ---------------------')
    cdd_filter_lungmask(thresh_r=10,
                        candis_dir=os.path.join(cfg.candi_dir, fold_name),
                        candis_lungmask_dir=os.path.join(cfg.candis_lungmask_dir, fold_name),
                        kaggle_lungmask_path=cfg.kaggle_lungmask_path,
                        lidc_lungmask_path=cfg.lidc_lungmask_path,
                        spie_lungmask_path=cfg.spie_lungmask_path,
                        stage2_lungmask_path=cfg.stage2_lungmask_path,
                        kaggle_vol_path=cfg.kaggle_vol_path,
                        lidc_vol_path=cfg.lidc_vol_path,
                        spie_vol_path=cfg.spie_vol_path,
                        stage2_vol_path=cfg.stage2_vol_path)

    print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
    print('---------------- Generate Candidates List ------------------')
    print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
    fold_idx = 9
    gen_list_one_fold(candi_dir=cfg.candis_lungmask_dir,
                      fold_name=fold_names[fold_idx],
                      candis_thresh=cfg.candis_thresh,
                      pkl_dir=cfg.pkl_dir)

